package cn.me.implement.threadpool;

import cn.me.implement.threadpool.reject.RejectPolicy;
import cn.me.implement.threadpool.reject.ThrowRejectPolicy;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * 自定义线程池
 * <p>
 * 在本实现中，将线程分为两种：
 * <ul>
 *     <li>核心线程：<strong>只要启动就一直运行，直到线程池 shutdown</strong></li>
 *     <li>支持线程：<strong>当一段时间没有任务就停止</strong></li>
 * </ul>
 */
public class JaneThreadPool {

    /**
     * 核心线程数量
     */
    private final int corePoolSize;

    /**
     * 最大线程数量
     */
    private final int maxPoolSize;

    /**
     * 如果支持线程在 {@link #keepAliveTime} 的单位时间内没有从任务队列中拿到任务，则支持线程停止
     */
    private final int keepAliveTime;

    /**
     * 时间的单位
     */
    private final TimeUnit timeUnit;

    /**
     * 任务队列
     */
    private final BlockingQueue<Runnable> taskQueue;

    /**
     * 当 总线程数等于最大线程数量 且 任务队列已满 时的拒绝策略
     */
    private final RejectPolicy rejectPolicy;

    /**
     * 线程工厂
     */
    private final ThreadFactory threadFactory;

    public JaneThreadPool(int corePoolSize, int maxPoolSize, int keepAliveTime, TimeUnit timeUnit,
            BlockingQueue<Runnable> taskQueue) {
        this(corePoolSize, maxPoolSize, keepAliveTime, timeUnit, taskQueue, new ThrowRejectPolicy());
    }

    public JaneThreadPool(int corePoolSize, int maxPoolSize, int keepAliveTime, TimeUnit timeUnit,
            BlockingQueue<Runnable> taskQueue, ThreadFactory threadFactory) {
        this(corePoolSize, maxPoolSize, keepAliveTime, timeUnit, taskQueue, new ThrowRejectPolicy(), threadFactory);
    }

    public JaneThreadPool(int corePoolSize, int maxPoolSize, int keepAliveTime, TimeUnit timeUnit,
            BlockingQueue<Runnable> taskQueue, RejectPolicy rejectPolicy) {
        this(corePoolSize, maxPoolSize, keepAliveTime, timeUnit, taskQueue, rejectPolicy, Thread::new);
    }

    public JaneThreadPool(int corePoolSize, int maxPoolSize, int keepAliveTime, TimeUnit timeUnit,
            BlockingQueue<Runnable> taskQueue, RejectPolicy rejectPolicy, ThreadFactory threadFactory) {
        this.corePoolSize = corePoolSize;
        this.maxPoolSize = maxPoolSize;
        this.keepAliveTime = keepAliveTime;
        this.timeUnit = timeUnit;
        this.taskQueue = taskQueue;
        this.rejectPolicy = rejectPolicy;
        this.threadFactory = threadFactory;
    }

    /**
     * 保存核心线程的集合
     */
    private final List<Thread> coreThreadPool = new ArrayList<>();

    /**
     * 保存支持线程的集合
     */
    private final List<Thread> supportThreadPool = new ArrayList<>();

    /**
     * 核心线程池的管程
     * <p>
     * 用于保证 <strong>获取线程池大小</strong>, <strong>新增线程</strong>, <strong>移除线程</strong> 的互斥性
     * <p>
     * 使用 synchronized 而不使用 ReentrantLock 是因为 synchronized 用起来更简单
     * <p>
     * TODO 在 forEach 之外使用 synchronized 可能有些太浪费性能了，之后可以想一想有没有其他解决方案
     */
    private final Object corePoolMonitor = new Object();

    /**
     * 支持线程池的管程，用于保证 <strong>获取线程池大小</strong>, <strong>新增线程</strong>, <strong>移除线程</strong> 的互斥性
     */
    private final Object supportPoolMonitor = new Object();

    /**
     * 任务队列的管程，用于保证 <strong>获取任务队列大小</strong> 和 <strong>清空队列</strong> 的互斥性
     */
    private final Object taskQueueMonitor = new Object();

    /**
     * 判断当前线程池是否处于 <strong>正在关闭/关闭</strong> 的状态
     */
    private final AtomicBoolean isShuttingDown = new AtomicBoolean(false);

    /**
     * 执行指定任务
     *
     * @param task 指定的任务
     */
    public void execute(Runnable task) {
        // 如果线程池处于 正在关闭/关闭 的状态，则直接拒绝任务
        if (isShuttingDown()) {
            rejectPolicy.reject(task, this);
        }

        // 如果 当前核心线程数 还没有到达 corePoolSize，则直接启动一个 核心线程 来处理任务
        synchronized (corePoolMonitor) {
            if (coreThreadPool.size() < corePoolSize) {
                startCoreThread(task);
                return;
            }
        }

        // 如果任务队列能放入新任务，则直接返回
        if (taskQueue.offer(task)) {
            return;
        }

        // 如果 当前总线程数 还没有到达 maxPoolSize，则直接启动一个 支持线程 来处理任务
        synchronized (corePoolMonitor) {
            synchronized (supportPoolMonitor) {
                if (coreThreadPool.size() + supportThreadPool.size() < maxPoolSize) {
                    startSupportThread(task);
                    return;
                }
            }
        }

        // 再看一下任务队列能否放入新任务，如果不能则触发拒绝策略
        if (!taskQueue.offer(task)) {
            rejectPolicy.reject(task, this);
        }
    }

    // 启动一个核心线程，它的第一个任务是 task
    private void startCoreThread(Runnable task) {
        CoreTask coreTask = new CoreTask(task);
        Thread coreThread = threadFactory.newThread(coreTask);
        coreTask.setThread(coreThread);
        coreThreadPool.add(coreThread);
        coreThread.start();
    }

    // 启动一个支持线程，它的第一个任务是 task
    private void startSupportThread(Runnable task) {
        SupportTask supportTask = new SupportTask(task);
        Thread supportThread = threadFactory.newThread(supportTask);
        supportTask.setThread(supportThread);
        supportThreadPool.add(supportThread);
        supportThread.start();
    }

    /**
     * 关闭线程池：
     * <ul>
     *     <li>给每个线程发送一个中断信号，等其执行完任务，就可以停止</li>
     *     <li>清理任务队列</li>
     * </ul>
     */
    public void shutdown() {
        doShutdown();
    }

    /**
     * 关闭线程池内部调用的方法
     * <p>
     * 这个方法仅在测试时使用
     *
     * @return 返回一个标志位表示是否是当前线程执行的 {@link #shutdown()} 导致线程池关闭
     */
    final boolean doShutdown() {
        // 如果线程池已经处于关闭状态，则直接返回，避免重复关闭线程池
        if (isShuttingDown() || !setIsShuttingDownToTrue()) {
            return false;
        }

        // 给每个线程发送中断信号
        synchronized (corePoolMonitor) {
            coreThreadPool.forEach(Thread::interrupt);
        }
        synchronized (supportPoolMonitor) {
            supportThreadPool.forEach(Thread::interrupt);
        }

        // 清理任务队列
        synchronized (taskQueueMonitor) {
            taskQueue.clear(); // help GC
        }

        return true;
    }

    /**
     * 移除任务队列中最旧的任务
     * <p>
     * 在拒绝策略中使用
     */
    public void pollOldestTask() {
        taskQueue.poll();
    }

    /**
     * 判断线程池是否处于 <strong>正在关闭/关闭</strong> 状态
     *
     * @return 线程池是否处于 <strong>正在关闭/关闭</strong> 状态
     */
    public boolean isShuttingDown() {
        return isShuttingDown.get();
    }

    /**
     * 将 {@link #isShuttingDown} 的状态从 {@code false} 设置为 {@code true}
     *
     * @return 是否设置成功
     */
    private boolean setIsShuttingDownToTrue() {
        return isShuttingDown.compareAndSet(false, true);
    }

    /**
     * 获取当前核心线程池中的线程数
     *
     * @return 当前核心线程池中的线程数
     */
    public int getCurrCoreThreadNum() {
        synchronized (corePoolMonitor) {
            return coreThreadPool.size();
        }
    }

    /**
     * 获取当前线程池中的所有线程数
     *
     * @return 当前线程池中的所有线程数
     */
    public int getCurrThreadNum() {
        synchronized (corePoolMonitor) {
            synchronized (supportPoolMonitor) {
                return coreThreadPool.size() + supportThreadPool.size();
            }
        }
    }

    /**
     * 获取当前任务队列中的任务数
     *
     * @return 当前任务队列中的任务数
     */
    public int getCurrTaskNum() {
        synchronized (taskQueueMonitor) {
            return taskQueue.size();
        }
    }

    /**
     * <h3>工作线程</h3>
     * 作为 <strong>核心线程</strong> 和 <strong>支持线程</strong> 的 <strong>抽象公共父类</strong>，保存一些公共字段
     * <p>
     * 这里使用了 <strong>模板方法模式</strong>，将子类需要执行的任务封装在 {@link WorkerTask#executeTask()} 方法中
     */
    private abstract static class WorkerTask implements Runnable {

        /**
         * 第一个任务，在新建线程时传递
         */
        private Runnable firstTask;

        /**
         * 工作线程在运行时绑定的线程
         */
        protected Thread thread;

        public WorkerTask(Runnable firstTask) {
            this.firstTask = firstTask;
        }

        @Override
        public final void run() {
            // 执行第一个任务，执行完毕清除对第一个任务的引用
            firstTask.run();
            firstTask = null; // help GC
            // 执行子类封装的处理任务逻辑
            executeTask();
            // 执行线程中断后的逻辑
            afterInterrupt();
        }

        /**
         * 处理任务：封装子类具体从任务队列中取出任务，并执行任务的逻辑
         */
        protected abstract void executeTask();

        /**
         * 在响应中断后执行的代码，一般用来 <strong>将本线程从线程池中移除</strong>
         */
        protected abstract void afterInterrupt();

        public void setThread(Thread thread) {
            this.thread = thread;
        }
    }

    /**
     * 核心线程
     */
    private class CoreTask extends WorkerTask {

        public CoreTask(Runnable firstTask) {
            super(firstTask);
        }

        @Override
        protected void executeTask() {
            while (Thread.currentThread().isInterrupted()) {
                Runnable task = null;
                try {
                    task = taskQueue.take();
                } catch (InterruptedException e) {
                    Thread.currentThread().interrupt();
                }
                if (task != null) {
                    task.run();
                }
            }
        }

        @Override
        protected void afterInterrupt() {
            synchronized (corePoolMonitor) {
                coreThreadPool.remove(thread);
            }
        }
    }

    /**
     * 支持线程
     */
    private class SupportTask extends WorkerTask {

        public SupportTask(Runnable firstTask) {
            super(firstTask);
        }

        @Override
        protected void executeTask() {
            while (Thread.currentThread().isInterrupted()) {
                Runnable task = null;
                try {
                    task = taskQueue.poll(keepAliveTime, timeUnit);
                } catch (InterruptedException e) {
                    Thread.currentThread().interrupt();
                }
                if (task == null) {
                    Thread.currentThread().interrupt();
                } else {
                    task.run();
                }
            }
        }

        @Override
        protected void afterInterrupt() {
            synchronized (supportPoolMonitor) {
                supportThreadPool.remove(thread);
            }
        }
    }
}
